# This m-script performs monte carlo simulation in Endogenous Treatment Models with Social Interactions (Lin and Vella 2024)
# Date: Jan 23, 2024
rm(list = ls())
ptm <- proc.time()
set.seed(1000)

library(GJRM)
library(VGAM)


n=250;
beta1=c(-1,1,1)
beta2=c(-1,1)
gamma=1
alpha=1
delta=1



library(mvtnorm)
nloop=200  ##Total number of replications
est=matrix(rep(0,9*nloop),nloop,9);
ATE=matrix(rep(0,2*nloop),nloop,2);
detectpv=matrix(rep(0,2*nloop),nloop,2);
for(l in 1:nloop){print(l)
  vc = diag(2)
  covs=-0.5
  vc[lower.tri(vc)] = covs
  vc[upper.tri(vc)] = vc[lower.tri(vc)]
  eps = rmvnorm(n, rep(0, 2), vc)
  
  
  
  x1=rnorm(n,mean=1)
  x2=rnorm(n)
  X1=cbind(1,x1)
  X2=cbind(1,x2)
  X=cbind(1,x1,x2)
  
  Friend=matrix(rep(0,n*n),n,n);
  NF = sample(0:10,n,TRUE);##Generating Friendship with maximum 10 friend.
  for(i in 1:n){
    location=sample(1:n,NF[i],TRUE);
    Friend[i,location]=1;
    Friend[i,i]=0;
  } 
  NF = rowSums(Friend);
  invNF =1/NF;#1 over number of friend
  invNF[which(!is.finite(invNF))]=0;
  
  t=10
  P0=rep(0,n);
  P1=rep(0,n);
  tot=100;
  reps=0;
  while (t > 1E-5&reps<tot){
    reps=reps+1;
    P1 = pnorm(X%*%beta1+alpha*Friend%*%P0*invNF);
    t=norm(P1-P0,"f")/sqrt(n);
    P0=P1;
  }
  
  S=as.matrix(Friend)%*%as.vector(P0)*as.vector(invNF)
  Y1 = (X%*%beta1+alpha*S+eps[,1]>0);##The treatment choice based on Cost-Benefit comparison
  Y0=1-Y1
  t=10
  PY0=rep(0,n);
  PY1=rep(0,n);
  tot=100;
  reps=0;
  while (t > 1E-5&reps<tot){
    reps=reps+1;
    IndexS=X%*%beta1+alpha*S;
    IndexO=X2%*%beta2+gamma*Y1+delta*Friend%*%PY0*invNF
    part1=(pnorm(IndexO)
           -pbinorm(IndexS,IndexO,cov12=covs))/(1-pnorm(IndexS));
    part2=Y1*(pbinorm(IndexS,IndexO,cov12=covs)/pnorm(IndexS)-(pnorm(IndexO)
                                                         -pbinorm(IndexS,IndexO,cov12=covs))/(1-pnorm(IndexS)));
    PY1 =part1+part2;
    t=norm(PY1-PY0,"f")/sqrt(n);
    PY0=PY1;
  }
  SY=as.matrix(Friend)%*%as.vector(PY0)*as.vector(invNF)
  
  Y2 = (X2%*%beta2+gamma*Y1+delta*SY+eps[,2]>0);##The treatment choice based on Cost-Benefit comparison
  AD=as.matrix(Friend)%*%as.vector(Y1)*as.vector(invNF)
  
  
  #####Calculate the true AME
  ME=rep(0,n)
  for(i in 1:n){
    if(Y1[i]==1){
      PO1=PY1;###Get P_O^1
      Y1C=Y1;Y1C[i]=0;
      tC=10
      PYC0=rep(0,n);
      PYC1=rep(0,n);
      totC=100;
      repsC=0;
      while (tC > 1E-5&repsC<totC){
        repsC=repsC+1;
        pp1=(pnorm(X2%*%beta2+gamma*Y1C+delta*Friend%*%PYC0*invNF)
             -pbinorm(X%*%beta1+alpha*S,X2%*%beta2+gamma*Y1C+delta*Friend%*%PYC0*invNF,cov12=covs))/(1-pnorm(X%*%beta1+alpha*S));
        pp2=Y1C*(pbinorm(X%*%beta1+alpha*S,X2%*%beta2+gamma*Y1C+delta*Friend%*%PYC0*invNF,cov12=covs)/pnorm(X%*%beta1+alpha*S)-(pnorm(X2%*%beta2+gamma*Y1C+delta*Friend%*%PYC0*invNF)
                     -pbinorm(X%*%beta1+alpha*S,X2%*%beta2+gamma*Y1C+delta*Friend%*%PYC0*invNF,cov12=covs))/(1-pnorm(X%*%beta1+alpha*S)));
        PYC1 = pp1+pp2;
        tC=norm(PYC1-PYC0,"f")/sqrt(n);
        PYC0=PYC1;
      }
      PO0=PYC0 ###Get P_O^0
    }
    else{   PO0=PY1;###Get P_O^0
    Y1CC=Y1;Y1CC[i]=1;
    tCC=10
    PYCC0=rep(0,n);
    PYCC1=rep(0,n);
    totCC=100;
    repsCC=0;
    while (tCC > 1E-5&repsCC<totCC){
      repsCC=repsCC+1;
      pp3=(pnorm(X2%*%beta2+gamma*Y1CC+delta*Friend%*%PYCC0*invNF)
           -pbinorm(X%*%beta1+alpha*S,X2%*%beta2+gamma*Y1CC+delta*Friend%*%PYCC0*invNF,cov12=covs))/(1-pnorm(X%*%beta1+alpha*S));
      pp4=Y1CC*(pbinorm(X%*%beta1+alpha*S,X2%*%beta2+gamma*Y1CC+delta*Friend%*%PYCC0*invNF,cov12=covs)/pnorm(X%*%beta1+alpha*S)-(pnorm(X2%*%beta2+gamma*Y1CC+delta*Friend%*%PYCC0*invNF)
                    -pbinorm(X%*%beta1+alpha*S,X2%*%beta2+gamma*Y1CC+delta*Friend%*%PYCC0*invNF,cov12=covs))/(1-pnorm(X%*%beta1+alpha*S)));
      PYCC1 = pp3+pp4;
      tCC=norm(PYCC1-PYCC0,"f")/sqrt(n);
      PYCC0=PYCC1;
    }
    PO1=PYCC0 ###Get P_O^1
    }
    SY1=as.matrix(Friend)%*%as.vector(PO1)*as.vector(invNF)
    SY0=as.matrix(Friend)%*%as.vector(PO0)*as.vector(invNF)
    
    ME[i]=pbinorm(sum(X[i,]*beta1)+alpha*S[i],sum(X2[i,]*beta2)+gamma+delta*SY1[i],cov12=covs)/P0[i]
    -(pnorm(sum(X2[i,]*beta2)+delta*SY0[i])
      -pbinorm(sum(X[i,]*beta1)+alpha*S[i],sum(X2[i,]*beta2)+delta*SY0[i],cov12=covs))/(1-P0[i])
  }
  AME=mean(ME)
  
  
  
  
  ###Estimation
  
  probit0 = glm(Y1~X,family=binomial(link="probit"));
  BNE=predict(probit0,type="response");
  BNEUpdate=rep(0,n);
  BNE=as.vector(BNE);
  BNEUpdate=as.vector(BNEUpdate);
  BNEY=BNE
  BNEYUpdate=BNEUpdate
  V=as.matrix(Friend)%*%as.vector(BNE)*as.vector(invNF)
  W=as.matrix(Friend)%*%as.vector(BNEY)*as.vector(invNF)
  datam=data.frame(Y1,Y2,x1,x2,V,W)
  
  treat.eq <- Y1 ~ x1+x2+V
  out.eq <- Y2 ~x2+Y1+W
  f.list <- list(treat.eq, out.eq)
  mr <- c("probit", "probit")
  
  
  
  err=10;
  err.tol=1E-3;
  npl.reps = 0;
  tot.reps = 100;
  quiet=FALSE;
  
  while(err>err.tol&npl.reps<tot.reps){
    npl.reps = npl.reps + 1;
    bipro=gjrm(f.list, data=datam, model="B", margins=mr)
    betasest=as.numeric(bipro$coefficients);
    sest=as.numeric(bipro$theta)
    BNEUpdate=pnorm(X%*%betasest[1:3]+betasest[4]*V);
    p1=(pnorm(X2%*%betasest[5:6]+betasest[7]*Y1+betasest[8]*W)
        -pbinorm(X%*%betasest[1:3]+betasest[4]*V,X2%*%betasest[5:6]+betasest[7]*Y1+betasest[8]*W,cov12=sest))/(1-BNEUpdate)
    p2=Y1*(pbinorm(X%*%betasest[1:3]+betasest[4]*V,X2%*%betasest[5:6]+betasest[7]*Y1+betasest[8]*W,cov12=sest)/BNEUpdate-(pnorm(X2%*%betasest[5:6]+betasest[7]*Y1+betasest[8]*W)
               -pbinorm(X%*%betasest[1:3]+betasest[4]*V,X2%*%betasest[5:6]+betasest[7]*Y1+betasest[8]*W,cov12=sest))/(1-BNEUpdate));
    BNEYUpdate=p1+p2;
    err=norm(as.matrix(cbind(BNEUpdate,BNEYUpdate))-as.matrix(cbind(BNE,BNEY)),"f")/sqrt(n);
    BNE=BNEUpdate;
    BNEY=BNEYUpdate
    datam$V=as.matrix(Friend)%*%as.vector(BNE)*as.vector(invNF);####Social interactions term updated using the new choice probabilities
    datam$W=as.matrix(Friend)%*%as.vector(BNEY)*as.vector(invNF);
    if(npl.reps==tot.reps & err>err.tol) {
      cat("The NPL algorithm did not converge. Try increase reps or reduce tolerance");
    }
    
  }
  
  
  
  est[l,]=as.numeric(c((bipro$coefficients)[1:8],bipro$theta))
  
  
  ###Estmated Average Marginal Effects
  EME=rep(0,n)
  for(i in 1:n){
    if(Y1[i]==1){
      POE1=BNEY;###Get P_O^1
      Y1C=Y1;Y1C[i]=0;
      tC=10
      PYC0=rep(0,n);
      PYC1=rep(0,n);
      totC=100;
      repsC=0;
      while (tC > 1E-5&repsC<totC){
        repsC=repsC+1;
        pc1=(pnorm(X2%*%betasest[5:6]+betasest[7]*Y1C+betasest[8]*Friend%*%PYC0*invNF)
             -pbinorm(X%*%betasest[1:3]+betasest[4]*V,X2%*%betasest[5:6]+betasest[7]*Y1C+betasest[8]*Friend%*%PYC0*invNF,cov12=sest))/(1-pnorm(X%*%betasest[1:3]+betasest[4]*V));
        pc2=Y1C*(pbinorm(X%*%betasest[1:3]+betasest[4]*V,X2%*%betasest[5:6]+betasest[7]*Y1C+betasest[8]*Friend%*%PYC0*invNF,cov12=sest)/pnorm(X%*%betasest[1:3]+betasest[4]*V)
                 -(pnorm(X2%*%betasest[5:6]+betasest[7]*Y1C+betasest[8]*Friend%*%PYC0*invNF)
                   -pbinorm(X%*%betasest[1:3]+betasest[4]*V,X2%*%betasest[5:6]+betasest[7]*Y1C+betasest[8]*Friend%*%PYC0*invNF,cov12=sest))/(1-pnorm(X%*%betasest[1:3]+betasest[4]*V)));
        PYC1 = pc1+pc2;
        tC=norm(PYC1-PYC0,"f")/sqrt(n);
        PYC0=PYC1;
      }
      POE0=PYC0 ###Get P_O^0
    }
    else{   POE0=BNEY;###Get P_O^0
    Y1CC=Y1;Y1CC[i]=1;
    tCC=10
    PYCC0=rep(0,n);
    PYCC1=rep(0,n);
    totCC=100;
    repsCC=0;
    while (tCC > 1E-5&repsCC<totCC){
      repsCC=repsCC+1;
      pc3=(pnorm(X2%*%betasest[5:6]+betasest[7]*Y1CC+betasest[8]*Friend%*%PYCC0*invNF)
           -pbinorm(X%*%betasest[1:3]+betasest[4]*V,X2%*%betasest[5:6]+betasest[7]*Y1CC+betasest[8]*Friend%*%PYCC0*invNF,cov12=sest))/(1-pnorm(X%*%betasest[1:3]+betasest[4]*V));
      pc4=Y1CC*(pbinorm(X%*%betasest[1:3]+betasest[4]*V,X2%*%betasest[5:6]+betasest[7]*Y1CC+betasest[8]*Friend%*%PYCC0*invNF,cov12=sest)/pnorm(X%*%betasest[1:3]+betasest[4]*V)-(pnorm(X2%*%betasest[5:6]+betasest[7]*Y1CC+betasest[8]*Friend%*%PYCC0*invNF)
                    -pbinorm(X%*%betasest[1:3]+betasest[4]*V,X2%*%betasest[5:6]+betasest[7]*Y1CC+betasest[8]*Friend%*%PYCC0*invNF,cov12=sest))/(1-pnorm(X%*%betasest[1:3]+betasest[4]*V)));
      PYCC1 = pc3+pc4;
      tCC=norm(PYCC1-PYCC0,"f")/sqrt(n);
      PYCC0=PYCC1;
    }
    POE1=PYCC0 ###Get P_O^1
    }
    SY1=as.matrix(Friend)%*%as.vector(POE1)*as.vector(invNF)
    SY0=as.matrix(Friend)%*%as.vector(POE0)*as.vector(invNF)
    
    EME[i]=pbinorm(sum(X[i,]*betasest[1:3])+betasest[4]*V[i],sum(X2[i,]*betasest[5:6])+betasest[7]+betasest[8]*SY1[i],cov12=sest)/BNEUpdate[i]
    -(pnorm(sum(X2[i,]*betasest[5:6])+betasest[8]*SY0[i])
      -pbinorm(sum(X[i,]*betasest[1:3])+betasest[4]*S[i],sum(X2[i,]*betasest[5:6])+betasest[8]*SY0[i],cov12=sest))/(1-BNEUpdate[i])
  }
  
  
  AEME=mean(EME);
  ATE[l,]=c(AME,AEME)
}

true=c(beta1,alpha,beta2,gamma,delta,covs)
mattrue=t(matrix(rep(true,nloop),9,nloop));

AB=round(colMeans(est)-true,digits=3);
MSE=round(colMeans((est-mattrue)^2),digits=3)


AB
MSE

round(c(colMeans(ATE),mean((ATE[,2]-ATE[,1])^2)),digits=3)


proc.time() - ptm
